import torch
import models
import numpy as np
from PIL import Image
import cv2
from data_cell import prepare_image_cv2
import os


def Tes():
    resume = './ckpt_multi/cell_256.pth'
    img_dir = r'H:\data\cell\RCF\imgs'
    result_path = './examples/0.png'

    model = models.resnet101(pretrained=False).cuda()
    model.eval()

    # resume..
    checkpoint = torch.load(resume)
    model.load_state_dict(checkpoint)

    names = os.listdir(img_dir)
    for name in names:
        img_path = os.path.join(img_dir, name)
        original_img = np.array(cv2.imread(img_path), dtype=np.float32)
        h, w, _ = original_img.shape

        img = prepare_image_cv2(original_img)
        img = torch.from_numpy(img).unsqueeze(0).cuda()

        outs = model(img, (h, w))
        result = outs[-1].squeeze().detach().cpu().numpy()

        result = (result * 255).astype(np.uint8)
        cv2.imshow("result", cv2.resize(result, (600, 600)))
        cv2.waitKey(1000)
        # Image.fromarray(result).save(result_path)


def val():
    resume = './ckpt_multi/cell_256.pth'
    img_dir = r'H:\data\cell\all\imgs'
    label_dir = r'H:\data\cell\all\gt_inv'


    model = models.resnet101(pretrained=False).cuda()
    model.eval()
    # resume..
    checkpoint = torch.load(resume)
    model.load_state_dict(checkpoint)
    names = os.listdir(img_dir)

    for name in names:
        img_path = os.path.join(img_dir, name)
        name_s = name.split(".")
        name_l = name_s[0] + ".jpg"
        lb_file = os.path.join(label_dir, name_l)
        lb_uint8 = cv2.imread(lb_file, cv2.IMREAD_GRAYSCALE)
        cv2.imshow("lb_uint8", cv2.resize(lb_uint8, (600, 600)))

        img_uint8 = cv2.imread(img_path)
        imgGray = cv2.cvtColor(img_uint8, cv2.COLOR_BGR2GRAY)
        original_img = np.array(img_uint8, dtype=np.float32)
        cv2.imshow("img_uint8", cv2.resize(img_uint8, (600, 600)))
        h, w, _ = original_img.shape

        img = prepare_image_cv2(original_img)
        img = torch.from_numpy(img).unsqueeze(0).cuda()

        outs = model(img, (h, w))
        result = outs[-1].squeeze().detach().cpu().numpy()

        result = (result * 255).astype(np.uint8)
        cv2.imshow("result", cv2.resize(result, (600, 600)))

        size_ori = lb_uint8.shape
        imgShow2 = np.empty((size_ori[0], size_ori[1], 3), np.uint8)
        imgShow2[:, :, 0] = imgGray
        imgShow2[:, :, 1] = np.maximum(imgGray, lb_uint8) * 0.9 + imgGray * 0.1
        imgShow2[:, :, 2] = np.maximum(imgGray, result) * 0.9 + imgGray * 0.1
        cv2.imshow("imgShow2", cv2.resize(imgShow2, (600, 600)))
        cv2.waitKey(1000)
        # Image.fromarray(result).save(result_path)


if __name__ == '__main__':
    val()


